from rest_framework import serializers
from .models import ChatMessage, Reference

class ReferenceSerializer(serializers.ModelSerializer):
    class Meta:
        model = Reference
        fields = ['title', 'content', 'source']

class MessageSerializer(serializers.ModelSerializer):
    message = serializers.CharField(write_only=True)
    references = ReferenceSerializer(many=True, read_only=True)
    context = serializers.ListField(child=serializers.DictField(), required=False)

    class Meta:
        model = ChatMessage
        fields = ['message_id', 'role', 'content', 'message', 'timestamp', 'context', 'references']
        read_only_fields = ['message_id', 'role', 'content', 'timestamp', 'references']

    def create(self, validated_data):
        message = validated_data.pop('message')
        context = validated_data.pop('context', [])
        user = validated_data.pop('user')
        
        # 创建用户消息
        user_message = ChatMessage.objects.create(
            user=user,
            role='user',
            content=message
        )

        # TODO: 调用AI服务获取回复
        ai_response = "这是AI的回复"  # 这里需要替换为实际的AI调用
        
        # 创建AI回复消息
        ai_message = ChatMessage.objects.create(
            user=user,
            role='assistant',
            content=ai_response
        )

        # 创建引用
        Reference.objects.create(
            message=ai_message,
            title="示例文档",
            content="引用内容",
            source="知识库"
        )

        return ai_message

class MessageListSerializer(MessageSerializer):
    class Meta(MessageSerializer.Meta):
        fields = ['message_id', 'role', 'content', 'timestamp', 'references']
        read_only_fields = fields

class ChatRequestSerializer(serializers.Serializer):
    message = serializers.CharField(max_length=2000)
    context = serializers.ListField(
        child=serializers.DictField(
            child=serializers.CharField(),
            allow_empty=False
        ),
        required=False
    )

class MessageIDsSerializer(serializers.Serializer):
    message_ids = serializers.ListField(
        child=serializers.CharField(),
        required=False
    )

class ExportRequestSerializer(MessageIDsSerializer):
    format = serializers.ChoiceField(choices=['pdf', 'word', 'markdown']) 